#include "tcp_sender.hh"

#include "tcp_config.hh"

#include <cassert>
#include <random>

// Dummy implementation of a TCP sender

// For Lab 3, please replace with a real implementation that passes the
// automated checks run by `make check_lab3`.

template <typename... Targs>
void DUMMY_CODE(Targs &&.../* unused */) {}

using namespace std;

//! \param[in] capacity the capacity of the outgoing byte stream
//! \param[in] retx_timeout the initial amount of time to wait before retransmitting the oldest outstanding segment
//! \param[in] fixed_isn the Initial Sequence Number to use, if set (otherwise uses a random ISN)
TCPSender::TCPSender(const size_t capacity, const uint16_t retx_timeout, const std::optional<WrappingInt32> fixed_isn)
    : _isn(fixed_isn.value_or(WrappingInt32{random_device()()}))
    , _initial_retransmission_timeout{retx_timeout}
    , _retransmission_timeout(retx_timeout)
    , _srtt(retx_timeout / 2)
    , _stream(capacity)
    , _window_size(1) {}

uint64_t TCPSender::bytes_in_flight() const {
    if (_ackno > _next_seqno) {
        return 0;
    }
    return _next_seqno - _ackno;
}

void TCPSender::fill_window() {
    //    printf("fill_window, window_size: %ld\n", _window_size);
    if (eof) {
        return;
    }
    uint64_t i = _next_seqno;
    uint64_t bytes_unacceptable = _ackno + _window_size;
    // printf("fill_window, _ackno: %lu\n", _ackno);
    while (i < bytes_unacceptable) {
        if (_window.count(i)) {
            // printf("_window.count(i): %ld\n", i);
            i += _window.at(i).length_in_sequence_space();
            if (_window.at(i).header().fin) {
                break;
            }
            if (_window.at(i).length_in_sequence_space() == 0) {
                // printf("_window.at(i).length_in_sequence_space() == 0???\n");
                break;
            }
            continue;
        }
        TCPSegment tcpSegment = TCPSegment();
        tcpSegment.header().seqno = wrap(i, _isn);
        tcpSegment.header().ack = _established;
        tcpSegment.header().win = _receiver_window_size;
        if (_established) {
            tcpSegment.header().ackno = _receiver_ackno.value();
        }
        if (i == 0) {
            // printf("set syn\n");
            tcpSegment.header().syn = true;
        } else if (!_stream.eof()) {
            size_t nread = bytes_unacceptable - i < TCPConfig::MAX_PAYLOAD_SIZE ? bytes_unacceptable - i
                                                                                : TCPConfig::MAX_PAYLOAD_SIZE;
            // printf("nread: %ld\n", nread);
            string s = _stream.read(nread);
            if (s.empty()) {
                break;
            }
            tcpSegment.setPayload(s);
        } else {
            tcpSegment.header().fin = true;
            eof = true;
        }
        if (_stream.eof() && tcpSegment.payload().size() < bytes_unacceptable - i) {
            // printf("set fin\n");
            tcpSegment.header().fin = true;
            eof = true;
        }
        // printf("tcp.seg.len: %ld\n", tcpSegment.length_in_sequence_space());
        insert2window(i, tcpSegment);
        i += tcpSegment.length_in_sequence_space();
        if (tcpSegment.header().fin) {
            break;
        }
    }
    _next_seqno = i;
    // printf("_segments_out.size: %ld\n", _segments_out.size());
}

void TCPSender::insert2window(uint64_t i, const TCPSegment &tcpSegment) {
    if (_window.empty()) {
        _retransmission_timer = _ms_ticks + _retransmission_timeout;
    }
    _window.insert(std::pair<uint64_t, TCPSegment>(i, tcpSegment));
    auto timer = RetransmissionTimer(i);
    _retransmission_map.insert(std::pair<uint64_t, RetransmissionTimer>(i, timer));
    _segments_out.push(tcpSegment);
}

//! \param ackno The remote receiver's ackno (Kacknowledgment number)
//! \param window_size The remote receiver's advertised window size
void TCPSender::ack_received(const WrappingInt32 ackno, const uint16_t window_size) {
    // printf("ack_received\n");
    if (window_size == 0) {
        _window_size = 1;
        _zero_window = true;
    } else {
        _window_size = window_size;
        _zero_window = false;
    }
    uint64_t absolute_ackno = unwrap(ackno, _isn, _ackno);
    // Impossible ackno (beyond next seqno) is ignored
    if (absolute_ackno > _next_seqno || absolute_ackno < _ackno) {
        return;
    }

    uint64_t i = _ackno;
    _ackno = absolute_ackno;
    _retransmission_timeout = _initial_retransmission_timeout;

    while (i < _ackno) {
        size_t len = _window.find(i)->second.length_in_sequence_space();
        //        printf("next i is %ld\n", i);
        if (len == 0) {
            assert(0);
        }
        if (i + len > _ackno) {
            _ackno = i;
            break;
        }
        _window.erase(i);
        _retransmission_map.erase(i);
        i += len;
        //  When an ACK is received that acknowledges new data, restart the
        //  retransmission timer so that it will expire after RTO seconds
        //  (for the current value of RTO).
        _retransmission_timer = _ms_ticks + _retransmission_timeout;
    }
    if (_ackno == _next_seqno) {
        _window.clear();
        _retransmission_map.clear();
    }
}

//! \param[in] ms_since_last_tick the number of milliseconds since the last call to this method
void TCPSender::tick(const size_t ms_since_last_tick) {
    // printf("ms_since_last_tick: %ld\n", ms_since_last_tick);
    _ms_ticks += ms_since_last_tick;
    if (_window.count(_ackno)) {
        //        printf("now: %ld, rtt: %d, _retransmission_timer: %ld\n",
        //               _ms_ticks,
        //               _retransmission_timeout,
        //               _retransmission_timer);
        if (_ms_ticks >= _retransmission_timer) {
            // printf("_retransmission_timer timeout!!\n");
            _retransmission_map.find(_ackno)->second.retransmit();  // reset ticks
            // When filling window, treat a '0' window size as equal to '1' but don't back off RTO
            if (!_zero_window) {
                _retransmission_timeout *= 2;
            }
            //            printf("resend: %s, fin: %d\n\n",
            //                   _window.find(_ackno)->second.payload().copy().c_str(),
            //                   _window.find(_ackno)->second.header().fin);
            _segments_out.push(_window.find(_ackno)->second);
            _retransmission_timer = _ms_ticks + _retransmission_timeout;
            // printf("_retransmission_timer: %ld\n", _retransmission_timer);
        }
    }
}

unsigned int TCPSender::consecutive_retransmissions() const {
    if (_retransmission_map.empty()) {
        return 0;
    }
    RetransmissionTimer timer = _retransmission_map.find(_ackno)->second;
    return timer.retransmissions_cnt();
}

void TCPSender::send_empty_segment(bool ack, WrappingInt32 ackno) {
    // printf("send_empty_segment\n");
    TCPSegment tcpSegment = TCPSegment();
    // When the keepalive timer reaches zero,
    // you send your peer a keepalive probe packet with no data in it
    // and the ACK flag turned on.
    tcpSegment.header().ack = ack;
    tcpSegment.header().ackno = ackno;
    tcpSegment.header().win = _receiver_window_size;
    if (ack) {
        _established = true;
        _receiver_ackno = ackno;
    }
    if (_next_seqno == 0) {
        // syn
        tcpSegment.header().syn = true;
        tcpSegment.header().seqno = _isn;
        insert2window(0, tcpSegment);
        // printf("send syn\n");
        _next_seqno += tcpSegment.length_in_sequence_space();
        return;
    }else{
        tcpSegment.header().seqno = wrap(_next_seqno, _isn);
    }
    _segments_out.push(tcpSegment);
    // printf(" _segments_out.size(): %ld\n", _segments_out.size());
}

void TCPSender::send_rst_segment() {
    // printf("send_rst_segment\n");
    TCPSegment tcpSegment = TCPSegment();
    tcpSegment.header().rst = true;
    tcpSegment.header().seqno = next_seqno();
    _segments_out.push(tcpSegment);
}
void TCPSender::set_receiver_window_size(size_t recv_win_size) { _receiver_window_size = recv_win_size; }
RetransmissionTimer::RetransmissionTimer(uint64_t next_seqno) {
    _next_seqno = next_seqno;
    _retransmissions_cnt = 0;
}

void RetransmissionTimer::retransmit() {
    // printf("retransmit: %d\n", _retransmissions_cnt);
    _retransmissions_cnt++;
}

unsigned int RetransmissionTimer::retransmissions_cnt() const { return _retransmissions_cnt; }
